import logging
import pprint

import cv2
import mxnet as mx
import mxnet.ndarray as nd
import numpy as np
import tqdm
from gluoncv.utils.metrics.voc_detection import VOC07MApMetric

from utils.block import FasterRCNN
from utils.config import config
from data.bbox.voc import VOCDetection
from models.defomableconvnets.bbox.bbox_transform import bbox_pred, clip_boxes
from models.defomableconvnets.config import config, update_config
from models.defomableconvnets.dcnresnet101 import DCNResnet101
from models.defomableconvnets.image import transform


def im_detect_bbox_aug(net, ctx_list, nms_threshold, im, scales, pixel_means, flip = False, threshold = 1e-3, viz = False):

    all_bboxes = []
    all_scores = []
    img_ori = im.copy()
    for scale_min, scale_max in scales:
        fscale = 1.0 * scale_min / min(img_ori.shape[:2])
        img_resized = cv2.resize(img_ori,(0,0),fx = fscale,fy = fscale)
        h,w,c = img_resized.shape
        h_padded = h if h %32==0 else h + 32 - h % 32
        w_padded = w if w %32==0 else w + 32 - w % 32
        img_padded = np.zeros(shape = (h_padded,w_padded,c), dtype = img_resized.dtype)
        img_padded[:h,:w,:] = img_resized
        img = transform(img_padded,pixel_means=pixel_means)
        im_info = nd.array([[h_padded,w_padded,1.0]], ctx= ctx[0])
        data = nd.array(img,ctx=ctx_list[0])

        rois, _, _, scores, bbox_deltas = net(data, im_info)
        scores = scores.reshape((1, -1, config.dataset.NUM_CLASSES)).softmax(axis = 2)
        bbox_deltas = bbox_deltas.reshape((1, -1, 4 * (2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)))

        rois = rois[:,1:].asnumpy()
        bbox_deltas = bbox_deltas[0].asnumpy()
        bbox = bbox_pred(rois, bbox_deltas)
        bbox = clip_boxes(bbox,data.shape[2:4])
        bbox /= fscale
        all_bboxes.append(bbox)
        all_scores.append(scores[0].asnumpy())

        # flip
        if flip:
            rois, _, _, scores, bbox_deltas = net(data[:,:,:,::-1],im_info)
            scores = scores.reshape((1, -1, config.dataset.NUM_CLASSES)).softmax(axis=2)
            bbox_deltas = bbox_deltas.reshape((1, -1, 4 * (2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)))

            rois = rois[:,1:].asnumpy()
            bbox_deltas = bbox_deltas[0].asnumpy()
            bbox = bbox_pred(rois, bbox_deltas)
            bbox = clip_boxes(bbox,data.shape[2:4])

            tmp = bbox[:, 0::4].copy()
            bbox[:, 0::4] =data.shape[3] - bbox[:, 2::4] #x0 = w - x0
            bbox[:, 2::4] =data.shape[3] - tmp # x1 = w -x1
            bbox /= fscale

            all_bboxes.append(bbox)
            all_scores.append(scores[0].asnumpy())

    all_bboxes = np.concatenate(all_bboxes,axis = 0)
    all_scores = np.concatenate(all_scores,axis = 0)
    pred_bboxes = []
    pred_scores = []
    pred_clsid = []
    for j in range(1, all_scores.shape[1]):
        cls_scores = all_scores[:, j, np.newaxis]
        cls_boxes = all_bboxes[:, 4:8] if config.CLASS_AGNOSTIC else all_bboxes[:, j * 4:(j + 1) * 4]
        cls_dets = np.hstack((cls_boxes, cls_scores))
        cls_dets = nd.contrib.box_nms(nd.array(cls_dets,ctx=mx.gpu(ctx_list[0].device_id)),
                                      overlap_thresh=nms_threshold, coord_start=0, score_index=4, id_index=-1,
                                      force_suppress=True, in_format='corner',
                                      out_format='corner').asnumpy()
        cls_dets = cls_dets[cls_dets[:, -1] > threshold, :]
        pred_bboxes.append(cls_dets[:,:4])
        pred_scores.append(cls_dets[:,4])
        pred_clsid.append(j * np.ones(shape=(cls_dets.shape[0],), dtype=np.int))
    pred_bboxes = np.concatenate(pred_bboxes,axis = 0)
    pred_scores = np.concatenate(pred_scores,axis = 0)
    pred_clsid = np.concatenate(pred_clsid,axis = 0)
    if viz:
        import gluoncv
        import matplotlib.pyplot as plt
        gluoncv.utils.viz.plot_bbox(img_ori[:,:,::-1], bboxes=pred_bboxes,scores=pred_scores,labels=pred_clsid,thresh=1e-5)
        plt.show()
    return pred_bboxes,pred_scores,pred_clsid


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    update_config("configs/voc/resnet_v1_101_voc0712_rfcn_dcn_end2end_ohem_one_gpu.yaml")
    msg = pprint.pformat(config)
    logging.info(msg)
    ctx = [mx.gpu(int(i)) for i in config.gpus.split(',')]
    backbone = DCNResnet101()
    net = FasterRCNN(config, backbone)

    means = np.tile(np.array(config.TRAIN.BBOX_MEANS), 2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)
    stds = np.tile(np.array(config.TRAIN.BBOX_STDS), 2 if config.CLASS_AGNOSTIC else config.dataset.NUM_CLASSES)
    args_params = nd.load("output/rfcn_voc-9-0.814870423711.params")
    params = net.collect_params()
    weight = args_params['rfcn_bbox_weight']
    bias = args_params['rfcn_bbox_bias']
    repeat = bias.shape[0] / means.shape[0]

    args_params['rfcn_bbox_weight'] = weight * mx.nd.repeat(mx.nd.array(stds), repeats=repeat).reshape(
        (bias.shape[0], 1, 1, 1))
    args_params['rfcn_bbox_bias'] = args_params['rfcn_bbox_bias'] * mx.nd.repeat(mx.nd.array(stds), repeats=repeat) + mx.nd.repeat(
        mx.nd.array(means), repeats=repeat)

    for key in params: # type: str
        if key.startswith(net.prefix):
            params[key]._load_init(args_params[key.replace(net.prefix, "")], ctx=mx.cpu())
        else:
            params[key]._load_init(args_params[key], ctx=mx.cpu())

    net.collect_params().reset_ctx(ctx)

    val_dataset = VOCDetection(root="/media/kohill/data/kohill/VOCdevkit", splits=((2007, 'test'),))
    val_metric_5 = VOC07MApMetric(iou_thresh=.5)
    for i in tqdm.tqdm(range(len(val_dataset))):
        img_path, gt_boxes = val_dataset.at_with_image_path(i)
        pred_bboxes, pred_scores, pred_clsid = im_detect_bbox_aug(net, nms_threshold=config.TEST.NMS,
                                                                  im=cv2.imread(img_path),  # bgr
                                                                  pixel_means=config.network.PIXEL_MEANS,
                                                                  scales=config.SCALES,
                                                                  ctx_list=ctx,
                                                                  flip=False,
                                                                  threshold=1e-3,
                                                                  viz=True
                                                                  )
        val_metric_5.update(pred_bboxes=pred_bboxes[np.newaxis],
                            pred_labels=pred_clsid[np.newaxis] - 1,
                            pred_scores=pred_scores[np.newaxis],
                            gt_bboxes=gt_boxes[np.newaxis, :, :4],
                            gt_labels=gt_boxes[np.newaxis, :, 4],
                            gt_difficults=gt_boxes[np.newaxis, :, 5])
    re = val_metric_5.get()
    print (re)